use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use anyhow::{Result, anyhow};

/// 内存池实现，用于高效管理和重用内存
pub struct MemoryPool {
    // 总容量
    capacity: usize,
    // 已使用容量
    used: AtomicUsize,
    // 内存块映射
    blocks: HashMap<usize, Vec<u8>>,
    // 下一个块ID
    next_id: AtomicUsize,
}

impl MemoryPool {
    /// 创建新的内存池
    pub fn new(capacity: usize) -> Self {
        Self {
            capacity,
            used: AtomicUsize::new(0),
            blocks: HashMap::new(),
            next_id: AtomicUsize::new(1),
        }
    }
    
    /// 分配内存块
    pub fn allocate(&mut self, size: usize) -> Result<usize> {
        let current_used = self.used.load(Ordering::Relaxed);
        if current_used + size > self.capacity {
            return Err(anyhow!("Memory pool out of capacity"));
        }
        
        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
        let block = vec![0u8; size];
        self.blocks.insert(id, block);
        self.used.fetch_add(size, Ordering::Relaxed);
        
        Ok(id)
    }
    
    /// 释放内存块
    pub fn deallocate(&mut self, id: usize) -> bool {
        if let Some(block) = self.blocks.remove(&id) {
            let size = block.len();
            self.used.fetch_sub(size, Ordering::Relaxed);
            true
        } else {
            false
        }
    }
    
    /// 获取内存块引用
    pub fn get(&self, id: usize) -> Option<&[u8]> {
        self.blocks.get(&id).map(|block| block.as_slice())
    }
    
    /// 获取内存块可变引用
    pub fn get_mut(&mut self, id: usize) -> Option<&mut [u8]> {
        self.blocks.get_mut(&id).map(|block| block.as_mut_slice())
    }
    
    /// 获取当前使用的内存大小
    pub fn used_size(&self) -> usize {
        self.used.load(Ordering::Relaxed)
    }
    
    /// 获取总容量
    pub fn capacity(&self) -> usize {
        self.capacity
    }
    
    /// 获取可用容量
    pub fn available(&self) -> usize {
        self.capacity - self.used.load(Ordering::Relaxed)
    }
    
    /// 重置内存池
    pub fn reset(&mut self) {
        self.blocks.clear();
        self.used.store(0, Ordering::Relaxed);
    }
    
    /// 调整内存池大小
    pub fn resize(&mut self, new_capacity: usize) -> Result<()> {
        let current_used = self.used.load(Ordering::Relaxed);
        if new_capacity < current_used {
            return Err(anyhow!("Cannot resize to smaller than current usage"));
        }
        
        self.capacity = new_capacity;
        Ok(())
    }
}